# saes.py  ——  极简 S-AES 实现（16-bit 分组，16-bit 密钥）
# 参考：Holden “A Simplified AES Algorithm”（含 S-Box、MixColumns 与 GF(2^4) 细节）:contentReference[oaicite:1]{index=1}

# --- 常量：GF(2^4) 用的本原多项式 x^4 + x + 1 -> 0b1_0011 = 0x13 ---
_IRR = 0x13  # 10011b

# --- S-Box / InvS-Box（按文献给出的 16 元置换表）---
# nibble -> nibble（4bit 输入/输出）
SBOX = [
    0x9, 0x4, 0xA, 0xB,
    0xD, 0x1, 0x8, 0x5,
    0x6, 0x2, 0x0, 0x3,
    0xC, 0xE, 0xF, 0x7
]
INV_SBOX = [0]*16
for i, v in enumerate(SBOX):
    INV_SBOX[v] = i

# --- GF(16) 乘法：4bit 与 4bit，相乘并对 x^4+x+1 约化 ---
def gf16_mul(a: int, b: int) -> int:
    a &= 0xF; b &= 0xF
    res = 0
    for _ in range(4):
        if b & 1:
            res ^= a
        b >>= 1
        a <<= 1
        if a & 0x10:   # 溢出 bit
            a ^= _IRR  # 约化
    return res & 0xF

# --- 工具：16bit 与 2x2 nibbles 列主序互转 ---
# 块按列填充：state = [[n0, n2],
#                      [n1, n3]]  （每个 n* 为 4bit）
def block_to_state(x: int):
    x &= 0xFFFF
    n0 = (x >> 12) & 0xF
    n1 = (x >> 8)  & 0xF
    n2 = (x >> 4)  & 0xF
    n3 = x & 0xF
    return [[n0, n2],
            [n1, n3]]

def state_to_block(s):
    n0, n2 = s[0]
    n1, n3 = s[1]
    return ((n0 & 0xF) << 12) | ((n1 & 0xF) << 8) | ((n2 & 0xF) << 4) | (n3 & 0xF)

# --- 轮操作 ---
def add_round_key(state, rk16: int):
    k = block_to_state(rk16)
    return [[state[r][c] ^ k[r][c] for c in range(2)] for r in range(2)]

def sub_nibbles(state):
    return [[SBOX[state[r][c]] for c in range(2)] for r in range(2)]

def inv_sub_nibbles(state):
    return [[INV_SBOX[state[r][c]] for c in range(2)] for r in range(2)]

def shift_rows(state):
    # 第0行不变；第1行循环左移1
    return [state[0][:], [state[1][1], state[1][0]]]

def inv_shift_rows(state):
    # 第1行循环右移1
    return [state[0][:], [state[1][1], state[1][0]]]

def mix_columns(state):
    # 乘以 [[1,4],[4,1]] in GF(16)（最后一轮不做 MC）:contentReference[oaicite:2]{index=2}
    a00, a01 = state[0]
    a10, a11 = state[1]
    b00 = (gf16_mul(1, a00) ^ gf16_mul(4, a10)) & 0xF
    b10 = (gf16_mul(4, a00) ^ gf16_mul(1, a10)) & 0xF
    b01 = (gf16_mul(1, a01) ^ gf16_mul(4, a11)) & 0xF
    b11 = (gf16_mul(4, a01) ^ gf16_mul(1, a11)) & 0xF
    return [[b00, b01], [b10, b11]]

def inv_mix_columns(state):
    # [[9,2],[2,9]] 是 [[1,4],[4,1]] 的逆（在 GF(16) 下）
    a00, a01 = state[0]
    a10, a11 = state[1]
    b00 = (gf16_mul(9, a00) ^ gf16_mul(2, a10)) & 0xF
    b10 = (gf16_mul(2, a00) ^ gf16_mul(9, a10)) & 0xF
    b01 = (gf16_mul(9, a01) ^ gf16_mul(2, a11)) & 0xF
    b11 = (gf16_mul(2, a01) ^ gf16_mul(9, a11)) & 0xF
    return [[b00, b01], [b10, b11]]

# --- 密钥扩展（W0..W5 -> K0,K1,K2三个16bit 轮密钥）---
# Rcon(1)=1000b(=x^3)，Rcon(2)=0011b(=x+1)，与教材/讲义一致:contentReference[oaicite:3]{index=3}
def _rot_nib(byte):
    hi = (byte >> 4) & 0xF
    lo = byte & 0xF
    return ((lo << 4) | hi) & 0xFF

def _sub_byte(byte):
    return ((SBOX[(byte >> 4) & 0xF] << 4) | SBOX[byte & 0xF]) & 0xFF

def key_expand(key16: int):
    key16 &= 0xFFFF
    w0 = (key16 >> 8) & 0xFF
    w1 = key16 & 0xFF

    # Rcon1 = 0x80 (1000b||0000b), Rcon2 = 0x30 (0011b||0000b)
    R1 = 0x80
    R2 = 0x30

    w2 = w0 ^ R1 ^ _sub_byte(_rot_nib(w1))
    w3 = w2 ^ w1
    w4 = w2 ^ R2 ^ _sub_byte(_rot_nib(w3))
    w5 = w4 ^ w3

    k0 = ((w0 << 8) | w1) & 0xFFFF
    k1 = ((w2 << 8) | w3) & 0xFFFF
    k2 = ((w4 << 8) | w5) & 0xFFFF
    return k0, k1, k2

# --- 单块加/解密 ---
def encrypt_block16(plain16: int, key16: int) -> int:
    k0, k1, k2 = key_expand(key16)
    s = block_to_state(plain16)
    s = add_round_key(s, k0)

    # Round 1
    s = sub_nibbles(s)
    s = shift_rows(s)
    s = mix_columns(s)
    s = add_round_key(s, k1)

    # Round 2（无 MC）
    s = sub_nibbles(s)
    s = shift_rows(s)
    s = add_round_key(s, k2)
    return state_to_block(s)

def decrypt_block16(cipher16: int, key16: int) -> int:
    k0, k1, k2 = key_expand(key16)
    s = block_to_state(cipher16)

    # 逆 Round 2
    s = add_round_key(s, k2)
    s = inv_shift_rows(s)
    s = inv_sub_nibbles(s)

    # 逆 Round 1
    s = add_round_key(s, k1)
    s = inv_mix_columns(s)
    s = inv_shift_rows(s)
    s = inv_sub_nibbles(s)

    s = add_round_key(s, k0)
    return state_to_block(s)
